import torch

# 笔记.md
import torchvision.datasets as dataset  # 内置了一些公开数据集
import torchvision.transforms as transforms  # 数据处理，数据增强库

train_data = dataset.MNIST(root='../笔记.md/mnist',
                           train=True,
                           transform=transforms.ToTensor(),
                           download=True)
test_data = dataset.MNIST(root='../笔记.md/mnist',
                          train=False,
                          transform=transforms.ToTensor(),
                          download=True)

# batch_size 取出一部分来训练，而不是全丢入
import torch.utils.data as data_utils

train_loader = data_utils.DataLoader(dataset=train_data,
                                     batch_size=64,
                                     shuffle=True)  # 打乱
test_loader = data_utils.DataLoader(dataset=test_data,
                                    batch_size=64,
                                    shuffle=True)  # 打乱


# net
class CNN(torch.nn.Module):

    def __init__(self):
        super(CNN, self).__init__()
        # 通过序列工具来构造网络
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, kernel_size=5, padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.fc = torch.nn.Linear(14 * 14 * 32, 10)

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size()[0], -1)
        out = self.fc(out)
        return out


cnn = CNN()
cnn = cnn.cuda()

# loss
loss_func = torch.nn.CrossEntropyLoss()  # 交叉熵损失函数

# optimizer
optimizer = torch.optim.Adam(cnn.parameters(), lr=0.01)

# training
for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        images = images.cuda()
        labels = labels.cuda()

        outputs = cnn(images)
        loss = loss_func(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print('epoch is {}, ite is {}/{}, loss is {}'.format(epoch + 1, i, len(train_data) // 64, loss.item()))

    # test
    loss_test = 0
    accuracy_test = 0
    for i, (images, labels) in enumerate(test_loader):
        images = images.cuda()
        labels = labels.cuda()

        outputs = cnn(images)
        # outputs = batchsize * cls
        loss_test += loss_func(outputs, labels)
        _, pred = outputs.max(1)
        accuracy_test += (pred == labels).sum().item()

    accuracy_test = accuracy_test / len(test_data)
    loss_test = loss_test / (len(test_data) // 64)

    print('epoch is {}, accuracy is {}, loss_test is {}'.format(epoch+1,accuracy_test,loss_test.item()))

# save
torch.save(cnn,'../model/model_k2.pkl')

# 另一个文件
# load

# inference
